# -*- coding: UTF-8 -*-

from llm.llm import llm
from utils.utils import jsonOutputParser, build_vue_chart
from agent.build_vue import *

# from tools.invoke_api import *
# from trace_agent.trace_agent import *
# from utils.utils import *
from pydantic import BaseModel

from langgraph.graph import StateGraph, START, END

from langchain.prompts import (
    SystemMessagePromptTemplate,
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from pyquery import PyQuery as pq

import json


class ReactAgentState(BaseModel):
    Query: str = ""
    HeaderJson: dict = {}
    OverviewJson: dict = {}
    BarChartJson: dict = []
    PieChartJson: dict = []
    BuildJson: dict = {}


def Header(state: ReactAgentState) -> ReactAgentState:
    print("提取Header")

    sys_tmpl = """
    角色设定：网络安全可视化工程师，专注于将复杂的安全数据转化为直观的可视化展示

    任务要求：
    1. 根据提供的网络安全事件摘要，生成一个简洁但信息丰富的标题
    2. 编写一段详细的描述，概括主要安全事件和风险
    3. 标题应反映最高风险级别的事件
    4. 描述应包含关键事件类型、时间范围和严重程度

    输出规范：
    - 标题不超过15个字
    - 描述不超过100字
    - 使用专业但易懂的安全术语

    返回格式(严格遵循JSON格式):
    ```json
    {{
        "title":"生成的标题",
        "desc":"生成的描述内容"
    }}
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm | jsonOutputParser()

    res = chain.invoke({"summary": state.Query})

    print("提取到Header", res)

    state.HeaderJson = json.loads(res)

    return state


def Overview(state: ReactAgentState) -> ReactAgentState:
    print("提取Overview")

    sys_tmpl = """
    角色设定：网络安全数据分析师，擅长从安全事件中提取关键指标

    任务要求：
    1. 从提供的安全事件摘要中提取3-5个最重要的指标
    2. 每个指标应包含清晰的标签和准确的值
    3. 优先提取以下类型的指标：
       - 事件总数
       - 最高风险级别
       - 主要攻击类型
       - 时间跨度
       - 受影响资产

    输出规范：
    - 每个指标标签不超过8个字
    - 值应保持原始数据的准确性
    - 使用标准化的安全术语
    
    返回格式(严格遵循JSON数组格式):
    ```json
    [
        {{
            "label": "指标名称",
            "value": "指标值"
        }}
    ]
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm | jsonOutputParser()

    data = chain.invoke({"summary": state.Query})

    josn = {"type": "overview", "data": json.loads(data)}

    state.OverviewJson = josn

    print("提取到Overview", josn)

    return state


def BarChart(state: ReactAgentState) -> ReactAgentState:
    print("生成柱状图数据")

    sys_tmpl = """
    角色设定：网络安全数据分析师，擅长将安全事件数据可视化

    任务要求：
    1. 分析提供的网络安全事件数据，提取适合柱状图展示的对比数据
    2. 可以按时间、攻击类型或风险级别进行对比
    3. X轴应为分类项，Y轴应为数值
    4. 确保数据点数量在3-8个之间

    输出规范：
    - X轴标签不超过10个字
    - 提供清晰的图表标题
    - 使用标准化的网络安全术语

    返回格式:
    ```json
    [
        {{
          "title": "柱状图标题1",
          "type": "barChart",
          "data": [
            {{"label": "分类名称", "value": "数值"}},
          ]
        }},
        {{
          "title": "柱状图标题2",
          "type": "barChart",
          "data": [
            {{"label": "分类名称", "value": "数值"}},
          ]
        }}
    ]
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm | jsonOutputParser()

    res = chain.invoke({"summary": state.Query})

    state.BarChartJson = json.loads(res)

    print("提取到柱状图", json.loads(res))

    return state


def PieChart(state: ReactAgentState) -> ReactAgentState:
    print("生成饼图数据")

    sys_tmpl = """
    角色设定：网络安全数据分析师，擅长将安全事件数据可视化

    任务要求：
    1. 分析提供的网络安全事件数据，提取适合饼图展示的分类数据
    2. 主要关注攻击类型分布或风险级别分布
    3. 每个分类项应包含名称和数量
    4. 确保分类合理，不超过6个类别

    输出规范：
    - 使用标准化的网络安全术语
    - 数量应为整数
    - 其他类别应合并为"其他"

    返回格式:
    ```json
    [
        {{
          "title": "饼图1",
          "type": "barChart",
          "data": [
            {{"label": "分类名称", "value": "数值"}},
          ]
        }},
        {{
          "title": "饼图2",
          "type": "barChart",
          "data": [
            {{"label": "分类名称", "value": "数值"}},
          ]
        }}
    ]
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm | jsonOutputParser()

    res = chain.invoke({"summary": state.Query})

    state.PieChartJson = json.loads(res)

    print("提取到饼图", json.loads(res))

    return state


def Build(state: ReactAgentState) -> ReactAgentState:
    print("生成完整json")

    state.BuildJson = {
        "header": state.HeaderJson,
        "components": [
            state.OverviewJson,
            state.BarChartJson,
            state.PieChartJson,
        ],
    }

    print("state.BuildJson", state.BuildJson)

    # 保存json文件
    with open("build.json", "w", encoding="utf-8") as f:
        json.dump(state.BuildJson, f, indent=4, ensure_ascii=False)

    return state


def Html(state: ReactAgentState) -> ReactAgentState:
    print("生成HTML")

    sys_tmpl = """
    角色设定
    你是一名资深全栈开发工程师，专注于创建响应式、现代化且视觉吸引力强的网页界面。你擅长将设计概念转化为功能完善的前端实现，特别注重用户体验和性能优化。

    任务要求
    将输入内容转换为一个完整的HTML页面，要求：

    设计原则
    采用现代简约设计风格，保持界面干净整洁
    使用合理的视觉层次结构引导用户注意力
    确保良好的可读性和可访问性

    技术实现
    CSS框架选择：

    优先使用Tailwind CSS 3.0+（推荐）

    或Bootstrap 5.2+（备选）

    可结合自定义CSS实现特殊效果

    数据可视化：
    使用Chart.js 4.0+实现交互式数据可视化
    支持多种图表类型（折线图、柱状图、饼图等）
    添加平滑的动画过渡效果
    实现响应式图表，随容器大小调整

    额外要求
    添加适当的微交互提升用户体验
    确保代码符合W3C标准
    优化页面加载性能（考虑懒加载、资源压缩等）

    输出格式
    请按照以下结构组织代码：
    完整的HTML5文档结构
    内联或外联的CSS样式

    JavaScript实现（优先使用ES6+语法）
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm

    res = chain.invoke({"summary": json.dumps(state.BuildJson)})

    # 保存html
    with open("build.html", "w", encoding="utf-8") as f:
        f.write(res.text())

    return state


def BuildVue(state: ReactAgentState) -> ReactAgentState:
    print("生成vue文件")

    web_json = [
        {
            "type": "header",
            "filename": build_vue(state.HeaderJson, "Header.vue"),
        },
        {
            "type": "overview",
            "filename": build_vue(state.OverviewJson, "Overview.vue"),
        },
    ]

    for component in state.PieChartJson:
        web_json.append(
            {
                "type": "pieChart",
                "filename": build_vue_chart(component),
            }
        )

    for component in state.BarChartJson:
        web_json.append(
            {
                "type": "barChart",
                "filename": build_vue_chart(component),
            }
        )

    # 保存json
    with open("./web/public/components.json", "w", encoding="utf-8") as f:
        json.dump(web_json, f, indent=4, ensure_ascii=False)

    return state


# 提取纯HTMl
def Extract_Pure_HTML(state: ReactAgentState) -> ReactAgentState:
    print("提取纯HTML")

    sys_tmpl = """
    你是一位专业的web前端开发工程师，擅长编写语义化的HTML结构。请根据以下要求将输入内容转换为纯HTML格式：

    1. 输出内容必须严格限定在<body>标签内，不包含任何<head>、<!DOCTYPE>或<html>标签
    2. 禁止包含任何JavaScript代码或CSS样式，div可以添加class
    3. 使用语义化的HTML标签（如<article>、<section>、<header>等）合理组织内容
    4. 每个逻辑段落或内容区块必须包裹在<div class="module"></div>容器中
    5. 根据内容类型适当使用以下HTML元素：
      - 文本内容使用<p>标签
      - 标题使用<h1>-<h6>层级结构
      - 列表使用<ul>/<ol>和<li>
      - 表格使用<table>及相关元素
      - 图片使用<img>并保留alt属性
      - 链接使用<a>并保留href属性
    6. 尽可能保留原始内容的完整信息，包括所有细节
    7. 如果内容包含数据或项目列表，请为每个项目创建单独的结构化区块
    8. 确保生成的HTML代码格式良好、缩进正确且通过W3C验证

    示例输出结构：
    <div class="module">
      <h3>标题</h3>
    </div>
    <div class="module">
      <h4>标题</h4>
      <ul>
        xxxx
      </ul>
    </div>
    ```
    """

    human_tmpl = "输入内容：\n{summary}"

    system_prompt = SystemMessagePromptTemplate.from_template(sys_tmpl)
    human_prompt = HumanMessagePromptTemplate.from_template(human_tmpl)
    chat_prompt = ChatPromptTemplate.from_messages([system_prompt, human_prompt])

    chain = chat_prompt | llm

    res = chain.invoke({"summary": state.Query})

    print("modules", res.text())

    # 创建 PyQuery 对象
    doc = pq(res.text())

    # 选择所有 class 为 module 的元素并获取它们的 HTML
    modules = doc(".module")

    print("modules", modules)

    vues = []

    for module in modules.items():
        vues.append(build_vue(module.html()))

    print("vues", vues)

    # 保存json
    with open("./web/public/components.json", "w", encoding="utf-8") as f:
        json.dump(vues, f, indent=4, ensure_ascii=False)

    return state


workflow = StateGraph(ReactAgentState)

# workflow.add_node("Header", Header)
# workflow.add_node("Overview", Overview)
# workflow.add_node("BarChart", BarChart)
# workflow.add_node("PieChart", PieChart)
# # workflow.add_node("Build", Build)
# # workflow.add_node("Html", Html)
# workflow.add_node("BuildVue", BuildVue)

# workflow.add_edge(START, "Header")
# workflow.add_edge("Header", "Overview")
# workflow.add_edge("Overview", "BarChart")
# workflow.add_edge("BarChart", "PieChart")
# workflow.add_edge("PieChart", "BuildVue")

# workflow.add_edge("BuildVue", END)
# # workflow.add_edge("Build", "Html")
# # workflow.add_edge("Html", END)

workflow.add_node("Extract_Pure_HTML", Extract_Pure_HTML)

workflow.add_edge(START, "Extract_Pure_HTML")
workflow.add_edge("Extract_Pure_HTML", END)

invoke_api_graph = workflow.compile()

agent_state = ReactAgentState()
agent_state.Query = """
### 网络安全摘要信息

#### 资产网络攻击统计相关数据

- **资产 ID**: `0fd4b9d96d6f5ecab6c6a8b6809f5ae3`
- **资产名称**: `taomingwei-node3-ubuntu`
- **总事件数量**: 13 个
- **告警时间范围**：
  - 第一条告警时间：2022 年 4 月 3 日 09:52:39
  - 最新告警时间：2025 年 4 月 8 日 09:12:25

#### 资产网络攻击类型分布

- **攻击 ID 及描述**：

  - T1564.001：通过命令行创建隐藏文件和目录，告警级别为 3。
  - T1078.003：通过 Grep 命令查找系统上包含 Password 关键字的文件，告警级别为 3。
  - T1082：可疑 Modprobe 文件事件，告警级别为 4。
  - T1059：
    - 检测到 Linux tmp 目录程序运行和检测到 Linux 可疑 shell 命令，告警级别分别为 4（第 3、5 条数据）。
    - Linux 受限 Shell 通过 Linux 二进制突破，告警级别为 3。
    - 检测到 Linux perl 反弹 shell2，威胁级别: 4
  - T1027.001 (填充二进制文件)
  - T1070.002 (清除 Linux 系统日志)

- **攻击阶段**：
  - `TA0007`（第 2 条数据）
  - `TA0002`（第 3、4、5 条数据，以及检测到 Linux perl 反弹 shell 事件）
  - `TA0004`（填充二进制文件）
  - `TA0005`（清除 Linux 系统日志）

#### 资产漏洞相关情况

- **检测方式**：
  - 隐藏文件目录检测：通过命令行创建隐藏文件和目录。
  - 敏感命令执行检测：使用 Grep 命令查找系统上包含 Password 关键字的文件。
  - 发现系统信息检测：检测 modprobe 配置文件中涉及内核模块的文件事件。
  - 检测到 Linux tmp 目录程序运行（第 3 条数据）
  - Linux 受限 Shell 通过 Linux 二进制突破（第 4 条数据）
  - 检测到 Linux 可疑 shell 命令（第 5 条数据）
  - 填充二进制文件: 检测通过 DD 指令填充二进制的行为。
  - 清除 Linux 系统日志: 黑客通常在入侵结束后对入侵痕迹进行清理，例如对所有可能记录其攻击行为的日志进行删除操作。

#### 资产风险相关情况

- **告警级别**：
  - 第 3、5 条数据：告警级别为 4。
  - 第 4 条数据：告警级别为 3。
  - 填充二进制文件: 级别 3 (第 6 条)
  - 清除 Linux 系统日志: 级别 3 (第 7 条)
  - 检测到 Linux 可疑命令: 级别 4 (第 8 条)

#### 资产安全防护状态

- **修复建议**：
  - 排查和结束相关可疑进程。
  - 排查和删除相关文件。
  - 排查和删除相关系统设置。
  - 排查和删除相关注册表项。
  - 填充二进制文件:
    1. 通过任务管理排查和结束相关可疑进程。
    2. 通过资源管理器排查和删除相关文件。
    3. 排查和删除相关可疑系统设置。
    4. 排查和删除相关可疑注册表项。
  - 清除 Linux 系统日志:
    1. 确认删除的日志文件是否可疑，排除正常的运维与文件更新操作。
    2. 确认操作可疑的情况下，进一步排查导致文件删除的进程名以及对应的用户名，必要时结束进程。
    3. 对可疑进程的后续操作进行排查。
  - 检测到 Linux 可疑命令:
    排查命令上下文分析命令是否恶意，若为恶意及时结束相关进程。

### 总结

资产 `taomingwei-node3-ubuntu` 遭遇了多次网络安全事件：

1. 检测到 Linux 可疑命令。
2. 检测到 Linux perl 反弹 shell。

建议立即排查相关网络连接并采取措施防止进一步的恶意活动。
"""

invokeRet = invoke_api_graph.invoke(agent_state, config={"recursion_limit": 100})
# print(invokeRet)
